package com.hiss.basic.starter.websocket.manager;

import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson.JSONObject;
import com.hiss.basic.starter.websocket.model.PrincipalModel;
import org.springframework.web.socket.WebSocketSession;

import java.lang.ref.WeakReference;
import java.security.Principal;
import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.stream.Collectors;

public class WebSocketSessionManager {

    // 用户 web socket session 容器
    private static final Map<String, Set<WeakReference<WebSocketSession>>> USER_SESSIONS_MAP = new ConcurrentHashMap<>();

    // 渠道 web socket session 容器
    private static final Map<String, Set<WeakReference<WebSocketSession>>> CHANNEL_SESSIONS_MAP = new ConcurrentHashMap<>();

    public static void open(WebSocketSession session) {
        PrincipalModel model = getPrincipalModel(session);
        if (ObjectUtil.isNotNull(model)) {
            String userId = model.getUserId();
            String channel = model.getChannel();
            USER_SESSIONS_MAP.computeIfAbsent(userId, k -> new CopyOnWriteArraySet<>()).add(new WeakReference<>(session));
            CHANNEL_SESSIONS_MAP.computeIfAbsent(channel, k -> new CopyOnWriteArraySet<>()).add(new WeakReference<>(session));
        }
    }

    public static void close(WebSocketSession session) {
        PrincipalModel model = getPrincipalModel(session);
        if (ObjectUtil.isNotNull(model)) {
            String userId = model.getUserId();
            String channel = model.getChannel();
            if (USER_SESSIONS_MAP.containsKey(userId)) {
                USER_SESSIONS_MAP.get(userId).remove(session);
                if (USER_SESSIONS_MAP.get(userId).isEmpty()) {
                    USER_SESSIONS_MAP.remove(userId);
                }
            }
            if (CHANNEL_SESSIONS_MAP.containsKey(channel)) {
                CHANNEL_SESSIONS_MAP.get(channel).remove(session);
                if (CHANNEL_SESSIONS_MAP.get(channel).isEmpty()) {
                    CHANNEL_SESSIONS_MAP.remove(channel);
                }
            }
        }
    }

    public static Collection<WebSocketSession> getSessions() {
        return USER_SESSIONS_MAP.values().stream().flatMap(Set::stream)
                .map(WeakReference::get).filter(ObjectUtil::isNotNull).collect(Collectors.toSet());
    }

    public static Collection<WebSocketSession> getUserSessions(String userId) {
        return USER_SESSIONS_MAP.getOrDefault(userId, Collections.emptySet()).stream()
                .map(WeakReference::get).filter(ObjectUtil::isNotNull).collect(Collectors.toSet());
    }

    public static Collection<WebSocketSession> getChannelSessions(String channel) {
        return CHANNEL_SESSIONS_MAP.getOrDefault(channel, Collections.emptySet()).stream()
                .map(WeakReference::get).filter(ObjectUtil::isNotNull).collect(Collectors.toSet());
    }

    public static PrincipalModel getPrincipalModel(WebSocketSession session) {
        if (ObjectUtil.isNotNull(session) && ObjectUtil.isNotNull(session.getPrincipal())) {
            Principal principal = session.getPrincipal();
            return JSONObject.parseObject(principal.getName(), PrincipalModel.class);
        }
        return null;
    }
}